"""
处理数据集数据

数据集中正负样本是分开的，所以每个文件就一句话

处理目标
归到一个文本文件中，两列，第一列为评论  第二列为标签 正样本1 负样本0

"""
import os
import random

neg_test_path = r"E:\workspace\【4-python-workspace】\LLM\04_finetune\03_BERT_Pre-training+Fine-tuning(Movie review sentiment classification)\aclImdb\test\neg/"
neg_train_path = r"E:\workspace\【4-python-workspace】\LLM\04_finetune\03_BERT_Pre-training+Fine-tuning(Movie review sentiment classification)\aclImdb\train\neg/"
pos_test_path = r"E:\workspace\【4-python-workspace】\LLM\04_finetune\03_BERT_Pre-training+Fine-tuning(Movie review sentiment classification)\aclImdb\test\pos/"
pos_train_path = r"E:\workspace\【4-python-workspace】\LLM\04_finetune\03_BERT_Pre-training+Fine-tuning(Movie review sentiment classification)\aclImdb\train\pos/"

pos_train_listdir = os.listdir(pos_train_path)
pos_test_listdir = os.listdir(pos_test_path)
neg_train_listdir = os.listdir(neg_train_path)
neg_test_listdir = os.listdir(neg_test_path)
# 获取每个评论字符串
pos_content = []
neg_content = []

for txt_name in pos_train_listdir:
    file_path = os.path.join(pos_train_path, txt_name)
    with open(file_path, "r", encoding="utf-8") as file:
        content = file.read()
        content.replace("\t", "")
    pos_content.append([content, 1])

for txt_name in pos_test_listdir:
    file_path = os.path.join(pos_test_path, txt_name)
    with open(file_path, "r", encoding="utf-8") as file:
        content = file.read()
        content.replace("\t", "")
    pos_content.append([content, 1])

for txt_name in neg_train_listdir:
    file_path = os.path.join(neg_train_path, txt_name)
    with open(file_path, "r", encoding="utf-8") as file:
        content = file.read()
        content.replace("\t", "")
    neg_content.append([content, 0])

for txt_name in neg_test_listdir:
    file_path = os.path.join(neg_test_path, txt_name)
    with open(file_path, "r", encoding="utf-8") as file:
        content = file.read()
        content.replace("\t", "")
    neg_content.append([content, 0])

# 训练集 验证集 测试集比例  8:1:1
train_ratio = 0.8
val_ratio = 0.1
test_ratio = 0.1

pos_len = len(pos_content)
neg_len = len(neg_content)

train_data = pos_content[:int(pos_len * train_ratio)] + neg_content[:int(neg_len * train_ratio)]
val_data = pos_content[int(pos_len * train_ratio):int(pos_len * (train_ratio + val_ratio))] + \
           neg_content[int(neg_len * train_ratio):int(neg_len * (train_ratio + val_ratio))]
test_data = pos_content[int(pos_len * (1 - test_ratio)):] + neg_content[int(neg_len * (1 - test_ratio)):]

# 打乱顺序
random.shuffle(train_data)
random.shuffle(val_data)
random.shuffle(test_data)

# # 格式化数据，并且写入文件
# # 打开文件，使用 'w' 模式表示写入
with open("./data/train.csv", "w", encoding="utf-8") as f:
    # 遍历字符串列表
    for data in train_data:
        # 将字符串和制表符以及数字写入文件，每个字符串一行
        f.write(f"{data[0]}\t\t{data[1]}\n")

with open("./data/val.csv", "w", encoding="utf-8") as f:
    # 遍历字符串列表
    for data in val_data:
        # 将字符串和制表符以及数字写入文件，每个字符串一行
        f.write(f"{data[0]}\t\t{data[1]}\n")

with open("./data/test.csv", "w", encoding="utf-8") as f:
    # 遍历字符串列表
    for data in test_data:
        # 将字符串和制表符以及数字写入文件，每个字符串一行
        f.write(f"{data[0]}\t\t{data[1]}\n")

print("数据处理结束！")
print("训练集个数：{}".format(len(train_data)))
print("验证集个数：{}".format(len(val_data)))
print("测试集个数：{}".format(len(test_data)))
